#include "thread_pool.hpp"
#include <iostream>
#include <atomic>
#include <queue>
#include <mutex>
#include <condition_variable>
/*
[生产者线程] -----> [共享队列] -----> [线程池]
     |                                    |
     |                                    |
     |                                    ↓
     |                            [工作线程1] [工作线程2] [工作线程3] [工作线程4]
     |                                    ↓
     |                            执行消费者任务1,2,3
     |
     +-> 持续生产数据 ---> 放入队列 ---> 线程池中的线程负责消费
*/
void TestThdPool()
{
    // 测试1：创建不同大小的线程池
    {
        thread_pool pool1(1); // 单线程
        thread_pool pool2(4); // 多线程
        thread_pool pool3;    // 默认线程数
        std::cout << "=== 线程池创建测试通过 ===" << std::endl;
    }

    // 测试2：基本任务提交和执行
    {
        thread_pool pool(2);
        // 简单任务
        pool.add_task([]
                      { std::cout << "简单任务执行" << std::endl; });

        // 带返回值的任务
        int result = 0;
        pool.add_task([&result]
                      { result = 42; });

        std::this_thread::sleep_for(std::chrono::milliseconds(100));
        std::cout << "=== 基本任务执行测试通过 ===" << std::endl;
    }

    // 测试3：并发提交任务（你现有的测试）
    {
        thread_pool pool(2);
        std::thread thd1([&pool]
                         {
            for (int i = 0; i < 10; i++) {
                auto thdId = std::this_thread::get_id();
                pool.add_task([thdId]{
                    std::cout << "线程1任务: " << thdId << std::endl;
                });
            } });

        std::thread thd2([&pool]
                         {
            for (int i = 0; i < 10; i++) {
                auto thdId = std::this_thread::get_id();
                pool.add_task([thdId]{
                    std::cout << "线程2任务: " << thdId << std::endl;
                });
            } });

        thd1.join();
        thd2.join();
        std::this_thread::sleep_for(std::chrono::seconds(1));
        std::cout << "=== 并发提交测试通过 ===" << std::endl;
    }

    // 测试4：压力测试
    {
        thread_pool pool(4);
        const int TASK_COUNT = 1000;
        std::atomic<int> counter{0};

        for (int i = 0; i < TASK_COUNT; i++)
        {
            pool.add_task([&counter]
                          {
                counter++;
                std::this_thread::sleep_for(std::chrono::microseconds(10)); });
        }

        std::this_thread::sleep_for(std::chrono::seconds(2));
        std::cout << "完成任务数: " << counter << "/" << TASK_COUNT << std::endl;
        std::cout << "=== 压力测试通过 ===" << std::endl;
    }

    // // 测试5：异常处理
    // {
    //     thread_pool pool(2);
    //     pool.add_task([]
    //                   { throw std::runtime_error("测试异常"); });

    //     std::this_thread::sleep_for(std::chrono::milliseconds(100));
    //     std::cout << "=== 异常处理测试通过 ===" << std::endl;
    // }

    // 测试6：生产者-消费者模式测试
    {
        thread_pool pool(4); // 使用4个工作线程
        std::queue<int> data_queue;
        std::mutex queue_mutex;
        std::condition_variable cv;
        std::atomic<bool> stop_flag{false};
        std::atomic<int> produced_count{0};
        std::atomic<int> consumed_count{0};

        // 启动生产者线程
        std::thread producer([&]()
                             {
                                 for (int i = 0; i < 100 && !stop_flag; ++i)
                                 {
                                     {
                                         std::lock_guard<std::mutex> lock(queue_mutex);
                                         data_queue.push(i);
                                         produced_count++;
                                     }
                                     cv.notify_one();                                            // 通知消费者
                                     std::this_thread::sleep_for(std::chrono::milliseconds(10)); // 模拟生产耗时
                                 }
                                 stop_flag = true;
                                 cv.notify_all(); // 确保所有消费者被唤醒
                             });

        // 添加消费者任务到线程池
        const int CONSUMER_COUNT = 10;
        for (int i = 0; i < CONSUMER_COUNT; ++i)
        {
            pool.add_task([&, consumer_id = i]()
                          {
                while (!stop_flag || !data_queue.empty()) {
                    std::unique_lock<std::mutex> lock(queue_mutex);
                    cv.wait(lock, [&]() { 
                        return !data_queue.empty() || stop_flag; 
                    });

                    if (!data_queue.empty()) {
                        int value = data_queue.front();
                        data_queue.pop();
                        consumed_count++;
                        lock.unlock();
                        
                        // 模拟处理数据
                        std::cout << "消费者 " << consumer_id 
                                << " 处理数据: " << value << std::endl;
                        std::this_thread::sleep_for(std::chrono::milliseconds(50));
                    }
                } });
        }

        producer.join();
        std::this_thread::sleep_for(std::chrono::seconds(2)); // 等待所有消费完成

        std::cout << "生产的数据量: " << produced_count << std::endl;
        std::cout << "消费的数据量: " << consumed_count << std::endl;
        std::cout << "=== 生产者-消费者测试通过 ===" << std::endl;
    }
}

int main()
{
    TestThdPool();
    return 0;
}